from __future__ import annotations

"""
数据库与会话管理
-----------------
职责：
- 根据配置构建异步数据库连接 URL（兼容 psycopg/asyncpg 写法）
- 初始化 SQLModel 元数据（首次启动建表）
- 提供 FastAPI 依赖的异步会话 `get_session`

注意：
- 生产环境推荐使用迁移工具（如 Alembic）管理表结构；当前项目通过 `create_all` 快速初始化
- `expire_on_commit=False` 以避免提交后对象被过期（访问属性触发再次查询）
"""

from typing import AsyncGenerator

from sqlalchemy import text
from sqlalchemy.ext.asyncio import (
    AsyncEngine,
    AsyncSession,
    async_sessionmaker,
    create_async_engine,
)
from sqlmodel import SQLModel

from .config import get_settings

settings = get_settings()


def _build_async_database_url(url: str) -> str:
    """将同步/不同驱动的 URL 规范化为异步可用形式。

    支持：
    - postgresql+psycopg:// → postgresql+psycopg_async://（psycopg3 async 驱动）
    - postgresql+asyncpg:// 保持不变
    - postgresql:// → postgresql+asyncpg://（默认使用 asyncpg）
    """
    if "+psycopg://" in url:
        return url.replace("+psycopg://", "+psycopg_async://")
    if "+asyncpg" in url:
        return url
    return url.replace("postgresql://", "postgresql+asyncpg://")


DATABASE_URL = _build_async_database_url(settings.database_url)

engine: AsyncEngine = create_async_engine(DATABASE_URL, echo=False, future=True)
SessionLocal = async_sessionmaker(engine, class_=AsyncSession, expire_on_commit=False)


async def init_db() -> None:
    """在应用启动时创建全部表结构（若不存在）。"""
    async with engine.begin() as conn:
        await conn.run_sync(SQLModel.metadata.create_all)
        await conn.exec_driver_sql(
            "ALTER TABLE groups ADD COLUMN IF NOT EXISTS connector_name VARCHAR(255) NOT NULL DEFAULT 'default_connector'"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE group_members ADD COLUMN IF NOT EXISTS connector_name VARCHAR(255) NOT NULL DEFAULT 'default_connector'"
        )
        await conn.exec_driver_sql(
            "UPDATE groups SET connector_name = 'default_connector' WHERE connector_name IS NULL"
        )
        await conn.exec_driver_sql(
            "UPDATE group_members SET connector_name = 'default_connector' WHERE connector_name IS NULL"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE groups ALTER COLUMN connector_name DROP DEFAULT"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE group_members ALTER COLUMN connector_name DROP DEFAULT"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE scripts ADD COLUMN IF NOT EXISTS hidden BOOLEAN NOT NULL DEFAULT false"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE script_subscriptions ADD COLUMN IF NOT EXISTS hidden BOOLEAN NOT NULL DEFAULT false"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE scripts ALTER COLUMN hidden SET DEFAULT false"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE script_subscriptions ALTER COLUMN hidden SET DEFAULT false"
        )

        await conn.exec_driver_sql(
            "ALTER TABLE ai_auto_decisions ADD COLUMN IF NOT EXISTS candles JSONB"
        )
        
        # 确保 orders 表的 client_order_id 有 UNIQUE 约束（防重复下单）
        await conn.exec_driver_sql(
            """
            DO $$
            BEGIN
                IF NOT EXISTS (
                    SELECT 1 FROM pg_constraint 
                    WHERE conname = 'orders_client_order_id_key'
                ) THEN
                    ALTER TABLE orders ADD CONSTRAINT orders_client_order_id_key UNIQUE (client_order_id);
                END IF;
            END $$;
            """
        )
        
        # 为 runs 表添加调度器配置字段（用于重启后恢复）
        await conn.exec_driver_sql(
            "ALTER TABLE runs ADD COLUMN IF NOT EXISTS exchange TEXT"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE runs ADD COLUMN IF NOT EXISTS symbol TEXT"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE runs ADD COLUMN IF NOT EXISTS account_name TEXT"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE runs ADD COLUMN IF NOT EXISTS risk_config JSONB"
        )
        await conn.exec_driver_sql(
            "ALTER TABLE runs ADD COLUMN IF NOT EXISTS interval_seconds INTEGER"
        )
        
        # 为 decisions 表添加 ts 索引（用于优化按时间排序的查询）
        await conn.exec_driver_sql(
            """
            DO $$
            BEGIN
                IF NOT EXISTS (
                    SELECT 1 FROM pg_indexes 
                    WHERE tablename = 'decisions' AND indexname = 'idx_decisions_ts'
                ) THEN
                    CREATE INDEX idx_decisions_ts ON decisions (ts DESC);
                END IF;
            END $$;
            """
        )


async def get_session() -> AsyncGenerator[AsyncSession, None]:
    """FastAPI 依赖：提供一次请求范围内的异步数据库会话。"""
    async with SessionLocal() as session:
        await session.execute(text("SET TIME ZONE 'Asia/Shanghai'"))
        yield session
